#!/usr/bin/python
import csv
from optparse import OptionParser
import numpy as np

from pyx import *

def get_input(filename):
	"""
	This function will read in the stratigraphy data using whichever input 
	system we choose to use.
	"""
	layerdict =  dict()
	layerfile = open(filename)
	layer_reader = csv.reader(layerfile)
	for layer in layer_reader:
		tmplayer = dict()
		tmplayer["topheight"] = float(layer[0])
		tmplayer["topfmt"] = layer[1]
		tmplayer["bottomheight"] = float(layer[2])
		if layer[4] != "none":
			tmplayer["comments"] = layer[4]
		else:
			tmplayer["comments"] = ""
		tmplayer["litho"] = layer[5]
		tmplayer["grainsize"] = int(layer[6])
		try:
			grainfile = open(layer[3]+".dat")
			grain_reader = csv.reader(grainfile)
			tmpgrain = dict()
			for grainsize in grainfile:
				tmpgrain[grainsize[0]] = grainsize[1]
			tmplayer["graindata"] = tmpgrain
		except:
			tmplayer["graindata"] = None
		layerdict[layer[3]] = tmplayer
	return layerdict


def graph_layers(layerdict, units):
	"""
	This method will use matplotlib to graph the stratigraphic colum.
	"""
	hatchstyles = {"covered":draw_covered, "sandstone":draw_sandstone}
	stratc = canvas.canvas()
	depth="Height"
	if units != None:
		depth = depth+" ("+units+")"
	stratg = stratc.insert(graph.graphxy(width=5, height=20, y=graph.axis.linear(title=depth), x=graph.axis.linear(painter=None)))
	for layer in layerdict.values():
		top = layer["topheight"]
		bottom = layer["bottomheight"]
		if layer["topfmt"] == "sine":
			stratg.plot(graph.data.function("y(x)=0.1*sin(x)+"+str(top), min=0, max=100/layer["grainsize"]))
		elif layer["topfmt"] == "straight":
			stratg.plot(graph.data.function("y(x)="+str(top), min=0, max=100/layer["grainsize"]))
		stratg.plot(graph.data.function("y(x)="+str(bottom), min=0, max=100/layer["grainsize"]))
		if layer["litho"] == "covered":
			hatchstyles[layer["litho"]](stratg, top, bottom, layer["grainsize"])
		if layer["litho"] == "sandstone":
			hatchstyles[layer["litho"]](stratg, top, bottom, layer["grainsize"])
	for name in layerdict.keys():
		namepos = stratg.pos(100, layerdict[name]["bottomheight"]+0.5*(layerdict[name]["topheight"]-layerdict[name]["bottomheight"]))
		add_text(stratc, name, layerdict[name]["comments"], namepos)
		#stratc.text(namepos[0], namepos[1], name)
	stratc.writeEPSfile("column")
	return 0

def draw_covered(gr, top, bottom, grainsize):
	slope = (top-bottom)/100
	gr.plot(graph.data.function("y(x)=x*"+str(slope)+"+"+str(bottom), min=0, max=100/grainsize))
	gr.plot(graph.data.function("y(x)=x*-"+str(slope)+"+"+str(top), min=0, max=100/grainsize), [graph.style.line()])

def draw_sandstone(gr, top, bottom, grainsize):
	height = bottom

def add_text(canv, name, comment, position):
	string = name + "\\" + "\\ "
	string = string + "\\scriptsize{" + comment + "}"
	canv.text(position[0], position[1], string, [text.parbox(3)])

if __name__ == "__main__":
	text.set(mode="latex")
	use = "stratcolumn.py [options] layer_file"
	parser = OptionParser(usage=use)
	parser.add_option("-u", "--units", dest="units", 
	help="Units for height of the column layers")
	(options, args) = parser.parse_args()
	if len(args) < 1:
		print "Insufficient arguments. Use --help for assistance"
		exit(1)
	datfilename = args[0]
	layers = get_input(datfilename)
	graph_layers(layers, options.units)
